package com.example.text2sql.agent;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationOutput;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.alibaba.dashscope.utils.Constants;
import com.example.text2sql.domain.Table;
import com.example.text2sql.domain.TableColumn;
import com.example.text2sql.utils.DatabaseUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class Text2Sql {


    public static GenerationParam getGenerationParam(List<Message> messages) {
        return GenerationParam.builder()
                .model("qwen-turbo")
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .messages(messages)
                .build();
    }

    public static Generation getGeneration() {
        return new Generation();
    }


    static {

        Constants.apiKey = "sk-****";
    }

    public String llm(String prompt) throws NoApiKeyException, InputRequiredException {
        Message sysMessage1 = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content(prompt.toString())
                .build();

        List<Message> messages = new ArrayList<>();
        messages.add(sysMessage1);

        GenerationParam param = getGenerationParam(messages);
        List<GenerationOutput.Choice> choices = getGeneration().call(param).getOutput().getChoices();
        GenerationOutput.Choice choice = choices.get(0);
        Message message = choice.getMessage();
        String content = message.getContent();
        return content;
    }

    public String call(String userQuery) throws NoApiKeyException, InputRequiredException {

        List<Table> tables = DatabaseUtils.getAllTablesInfo();

        StringBuilder prompt = new StringBuilder();
        prompt.append("根据以下数据库表结构信息，将自然语言查询转换为 SQL 查询：\n\n");
        prompt.append("如果没有特殊要求, 不要自己发挥处理任何逻辑字段, 如：is_deleted。\n\n");
        prompt.append("回答结果中只需要返回我 SQL 内容, 非 md 格式。\n\n");
        for (Table table : tables) {
            prompt.append("表名: ").append(table.getName()).append("\n");
            prompt.append("描述: ").append(table.getRemarks()).append("\n");
            List<TableColumn> columns = DatabaseUtils.getTableColumnsInfo(table.getName());
            for (TableColumn column : columns) {
                prompt.append("字段名: ").append(column.getName()).append(", 类型: ").append(column.getDataType()).append(", 描述: ").append(column.getRemarks()).append("\n");
            }
            prompt.append("\n");
        }
        prompt.append("用户查询: ").append(userQuery).append("\n\n");


        String content = llm(prompt.toString());
        content = content.replaceAll("```sql", "");
        content = content.replaceAll("```", "");

        System.err.println("生成 SQL :" + content);

        List<Map<String, Object>> result = DatabaseUtils.executeSelectQuery(content);

        String resultStr = "";

        for (Map<String, Object> row : result) {
            for (Map.Entry<String, Object> entry : row.entrySet()) {
                resultStr += entry.getKey() + ": " + entry.getValue();
            }
        }


        StringBuilder prompt2 = new StringBuilder();
        prompt2.append("请你根据用户的意图, 和查询结果答复他的问题：\n\n");
        prompt2.append("查询结果：" + resultStr);
        prompt2.append("原始意图：" + userQuery);

        return llm(prompt2.toString());
    }
}
